{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segmented deformable mirrors\n",
    "\n",
    "We will use segmented deformable mirrors and simulate the PSFs that result from segment pistons and tilts. We will compare this functionality against Poppy, another optical propagation package.\n",
    "\n",
    "First we'll import all packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.units as u\n",
    "import hcipy\n",
    "import poppy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters - these will depend on the aperture function you use\n",
    "PUP_DIAMETER = 0.019725   # m\n",
    "GAPSIZE = 90e-6  # m\n",
    "FLATTOFLAT = PUP_DIAMETER/7.   # m\n",
    "\n",
    "# these will not\n",
    "num_pix = 1024\n",
    "wavelength = 638e-9\n",
    "num_airy = 20\n",
    "sampling = 4\n",
    "norm = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate the segmented mirrors\n",
    "\n",
    "### HCIPy SM: `hsm`\n",
    "\n",
    "We need to generate a pupil grid for the aperture, and a focal grid and propagator for the focal plane images after the DM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HCIPy grids and propagator\n",
    "pupil_grid = hcipy.make_pupil_grid(dims=num_pix, diameter=PUP_DIAMETER)\n",
    "focal_grid = hcipy.make_uniform_grid([2*num_airy*sampling]*2, 2*num_airy * wavelength / PUP_DIAMETER)\n",
    "prop = hcipy.FraunhoferPropagator(pupil_grid, focal_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate a segmented aperture for the segmented mirror. For convenience, we'll use the HiCAT pupil without spiders. We'll use supersampling to better resolve the segment gaps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aper, segments = hcipy.make_hicat_aperture(normalized=norm, with_spiders=False, return_segments=True)\n",
    "aper = hcipy.evaluate_supersampled(aper, pupil_grid, 1)\n",
    "segments = hcipy.evaluate_supersampled(segments, pupil_grid, 1)\n",
    "\n",
    "plt.title('HCIPy aperture')\n",
    "hcipy.imshow_field(aper)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we make the segmented mirror. In order to be able to apply the SM to a plane, that plane needs to be a `Wavefront`, which combines a `Field` - here the aperture - with a wavelength, here `wavelength`.\n",
    "\n",
    "In this example here, since the SM doesn't have any extra effects on the pupil since it's completely flat still, we don't actually have to apply the SM, although of course we could."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the segmented mirror\n",
    "hsm = hcipy.SegmentedDeformableMirror(segments)\n",
    "\n",
    "# Make a pupil plane wavefront from aperture\n",
    "wf = hcipy.Wavefront(aper, wavelength)\n",
    "\n",
    "# Apply SM if you want to\n",
    "wf = hsm(wf)\n",
    "\n",
    "plt.title('Wavefront intensity at HCIPy SM')\n",
    "hcipy.imshow_field(wf.intensity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poppy SM: `psm`\n",
    "\n",
    "We'll do the same for Poppy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psm = poppy.dms.HexSegmentedDeformableMirror(name='Poppy SM',\n",
    "                                             rings=3,\n",
    "                                             flattoflat=FLATTOFLAT*u.m,\n",
    "                                             gap=GAPSIZE*u.m,\n",
    "                                             center=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the transmission and phase of the poppy sm\n",
    "plt.figure(figsize=(16, 8))\n",
    "psm.display(what='both')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create reference images\n",
    "\n",
    "### HCIPy reference image\n",
    "\n",
    "We need to apply the SM to the wavefront in the pupil plane and then propagate it to the image plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply SM to pupil plane wf\n",
    "wf_sm = hsm(wf)\n",
    "\n",
    "# Propagate from SM to image plane\n",
    "im_ref_hc = prop(wf_sm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display intensity and phase in image plane\n",
    "plt.figure(figsize=(18, 9))\n",
    "plt.suptitle('Image plane after HCIPy SM')\n",
    "\n",
    "# Get normalization factor for HCIPy reference image\n",
    "norm_hc = np.max(im_ref_hc.intensity)\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "hcipy.imshow_field(np.log10(im_ref_hc.intensity/norm_hc))\n",
    "plt.title('Intensity')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "hcipy.imshow_field(im_ref_hc.phase, cmap='RdBu')\n",
    "plt.title('Phase')\n",
    "plt.colorbar()\n",
    "\n",
    "print('HCIPy PSF shape: {}'.format(im_ref_hc.intensity.shaped.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poppy reference image\n",
    "\n",
    "For the Poppy propagation, we need to make an optical system of which we then calculate the PSF.  \n",
    "\n",
    "I will try to match the image resolution and size of the HCIPy image. I first adjust the `pixelscale` and `fov_arcsec` such that their ratio works and then I add a tweak factor `fac` to scale it to the HCIPy image. I also set `oversample` to something that matches the HCIPy sampling (it's close enough). I keep reusing these numbers and the tweak factor later on in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make an optical system with the Poppy SM and a detector\n",
    "psm.flatten()\n",
    "osys = poppy.OpticalSystem()\n",
    "osys.add_pupil(psm)\n",
    "fac = 5300\n",
    "pxscle = 10 * np.degrees(wavelength / PUP_DIAMETER) * 3600.0 / sampling * 0.97\n",
    "fovarc = pxscle * 160 / 10\n",
    "osys.add_detector(pixelscale=pxscle, fov_arcsec=fovarc, oversample=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the PSF\n",
    "psf = osys.calc_psf(wavelength)\n",
    "plt.figure(figsize=(10, 10))\n",
    "poppy.display_psf(psf, vmin=1e-9, vmax=0.1)\n",
    "\n",
    "# Get the PSF as an array\n",
    "im_ref_pop = psf[0].data\n",
    "print('Poppy PSF shape: {}'.format(im_ref_pop.shape))\n",
    "\n",
    "# Get normalization from Poppy reference image\n",
    "norm_pop = np.max(im_ref_pop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Both reference images side-by-side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,10))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "hcipy.imshow_field(np.log10(im_ref_hc.intensity / norm_hc))\n",
    "plt.title('HCIPy reference PSF')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(np.log10(im_ref_pop / norm_pop), origin='lower')\n",
    "plt.title('Poppy reference PSF')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_dif = im_ref_pop / norm_pop - im_ref_hc.intensity.shaped / norm_hc\n",
    "\n",
    "plt.figure(figsize=(20, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(ref_dif, origin='lower')\n",
    "plt.title('Full image')\n",
    "plt.colorbar()\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(ref_dif[60:100,60:100], origin='lower')\n",
    "plt.title('Zoomed in')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Applying aberrations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function from rad of phase to m OPD\n",
    "def aber_to_opd(aber_rad, wavelength):\n",
    "    aber_m = aber_rad * wavelength / (2 * np.pi)\n",
    "    return aber_m\n",
    "    \n",
    "aber_rad = 4.0\n",
    "\n",
    "print('Aberration: {} rad'.format(aber_rad))\n",
    "print('Aberration: {} m'.format(aber_to_opd(aber_rad, wavelength)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flatten both SMs just to be sure\n",
    "hsm.flatten()\n",
    "psm.flatten()\n",
    "\n",
    "poppy_to_hcipy_index = {\n",
    "    1: 2, 2: 1, 3: 0, 4: 5, 5: 4, 6: 3,\n",
    "    7: 10, 8: 9, 9: 8, 10: 7, 11: 6, 12: 17, 13: 16, 14: 15, 15: 14, 16: 13, 17: 12, 18: 11,\n",
    "    19: 24, 20: 23, 21: 22, 22: 21, 23: 20, 24: 19, 25: 18,\n",
    "    26: 35, 27: 34, 28: 33, 29: 32, 30: 31, 31: 30, 32: 29, 33: 28, 34: 27, 35: 26, 36: 25}\n",
    "\n",
    "for i in [35, 25]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], aber_to_opd(aber_rad, wavelength) / 2, 0, 0)\n",
    "    psm.set_actuator(i, aber_to_opd(aber_rad, wavelength) * u.m, 0, 0)\n",
    "    \n",
    "# Display both segmented mirrors in OPD\n",
    "\n",
    "# HCIPy\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.title('Phase for HCIPy SM')\n",
    "hcipy.imshow_field(hsm.surface * 2, mask=aper, cmap='RdBu_r', vmin=-5e-7, vmax=5e-7)\n",
    "plt.colorbar()\n",
    "\n",
    "# Poppy\n",
    "plt.figure(figsize=(8,8))\n",
    "psm.display(what='opd')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show focal plane images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### HCIPy\n",
    "# Apply SM to pupil plane wf\n",
    "wf_fp_pistoned = hsm(wf)\n",
    "\n",
    "# Propagate from SM to image plane\n",
    "im_pistoned_hc = prop(wf_fp_pistoned)\n",
    "\n",
    "### Poppy\n",
    "# Make an optical system with the Poppy SM and a detector\n",
    "osys = poppy.OpticalSystem()\n",
    "osys.add_pupil(psm)\n",
    "pxscle = 0.0031*fac      # I'm tweaking pixelscale and fov_arcsec to match the HCIPy image\n",
    "fovarc = 0.05*fac\n",
    "osys.add_detector(pixelscale=pxscle, fov_arcsec=fovarc, oversample=10)\n",
    "\n",
    "# Calculate the PSF\n",
    "psf = osys.calc_psf(wavelength)\n",
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "# Get the PSF as an array\n",
    "im_pistoned_pop = psf[0].data\n",
    "\n",
    "### Display intensity of both cases image plane\n",
    "plt.figure(figsize=(18, 9))\n",
    "plt.suptitle('Image plane after SM for $\\phi$ = ' + str(aber_rad) + ' rad')\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "hcipy.imshow_field(np.log10(im_pistoned_hc.intensity / norm_hc))\n",
    "plt.title('HCIPy pistoned pair')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(np.log10(im_pistoned_pop / norm_pop), origin='lower')\n",
    "plt.title('Poppy pistoned pair')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image degradation as function of rms piston errors\n",
    "\n",
    "We will plot the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Aberration range\n",
    "aber_array = np.linspace(0, 2*np.pi, 50, True)\n",
    "print('Aber in rad: \\n{}'.format(aber_array))\n",
    "print('Aber in m: \\n{}'.format(aber_to_opd(aber_array, wavelength)))\n",
    "\n",
    "# Apply pistons\n",
    "hc_ims = []\n",
    "pop_ims = []\n",
    "for aber_rad in aber_array:\n",
    "    # Flatten both SMs\n",
    "    hsm.flatten()\n",
    "    psm.flatten()\n",
    "\n",
    "    # HCIPy\n",
    "    for i in [34, 25]:\n",
    "        opd = aber_to_opd(aber_rad, wavelength)\n",
    "        hsm.set_segment_actuators(poppy_to_hcipy_index[i], opd / 2, 0, 0)   # hcipy uses SURFACE, not OPD\n",
    "        psm.set_actuator(i, opd * u.m, 0, 0)\n",
    "\n",
    "    # Propagate to image plane\n",
    "    \n",
    "    # HCIPy:\n",
    "    # Propagate from pupil plane through SM to image plane\n",
    "    im_pistoned_hc = prop(hsm(wf))\n",
    "\n",
    "    # Poppy:\n",
    "    # Make an optical system with the Poppy SM and a detector\n",
    "    osys = poppy.OpticalSystem()\n",
    "    osys.add_pupil(psm)\n",
    "    pxscle = 0.0031 * fac  # I'm tweaking pixelscale and fov_arcsec to match the HCIPy image\n",
    "    fovarc = 0.05 * fac\n",
    "    osys.add_detector(pixelscale=pxscle, fov_arcsec=fovarc, oversample=10)\n",
    "\n",
    "    # Calculate the PSF\n",
    "    psf = osys.calc_psf(wavelength)\n",
    "\n",
    "    # Get the PSF as an array\n",
    "    im_pistoned_pop = psf[0].data\n",
    "\n",
    "    hc_ims.append(im_pistoned_hc.intensity.shaped / im_pistoned_hc.intensity.max())\n",
    "    pop_ims.append(im_pistoned_pop / im_pistoned_pop.max())\n",
    "    \n",
    "hc_ims = np.array(hc_ims)\n",
    "pop_ims = np.array(pop_ims)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Quantify with image sums\n",
    "sum_hc = np.sum(hc_ims, axis=(1,2))\n",
    "sum_hc /= sum_hc[0]\n",
    "sum_pop = np.sum(pop_ims, axis=(1,2))\n",
    "sum_pop /= sum_pop[0]\n",
    "\n",
    "plt.suptitle('Image degradation of SMs')\n",
    "plt.plot(aber_array, sum_hc, '-', label='HCIPy SM')\n",
    "plt.plot(aber_array, sum_pop, '--', label='Poppy SM')\n",
    "plt.xlabel('phase aberration (rad)')\n",
    "plt.ylabel('image sum')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A mix of piston, tip and tilt (PTT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aber_rad_tt = 500e-6\n",
    "aber_rad_p = 1.8\n",
    "\n",
    "opd_piston = aber_to_opd(aber_rad_p, wavelength)\n",
    "\n",
    "### Put aberrations on both SMs\n",
    "# Flatten both SMs\n",
    "hsm.flatten()\n",
    "psm.flatten()\n",
    "\n",
    "## PISTON\n",
    "for i in [19, 28, 23, 16]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], opd_piston / 2, 0, 0)\n",
    "    psm.set_actuator(i, opd_piston * u.m, 0, 0)\n",
    "    \n",
    "for i in [3, 35, 30, 8]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], 0.5 * opd_piston / 2, 0, 0)\n",
    "    psm.set_actuator(i, 0.5 * opd_piston * u.m, 0, 0)\n",
    "    \n",
    "for i in [14, 18, 1, 32, 12]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], 0.3 * opd_piston / 2, 0, 0)\n",
    "    psm.set_actuator(i, 0.3 * opd_piston * u.m, 0, 0)\n",
    "    \n",
    "## TIP and TILT\n",
    "for i in [2, 5, 11, 15, 22]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], 0, aber_rad_tt / 2, 0.3 * aber_rad_tt / 2)\n",
    "    psm.set_actuator(i, 0, aber_rad_tt, 0.3 * aber_rad_tt)\n",
    "    \n",
    "for i in [4, 6, 36]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], 0, aber_rad_tt / 2, 0)\n",
    "    psm.set_actuator(i, 0, aber_rad_tt, 0)\n",
    "    \n",
    "for i in [34, 31, 7]:\n",
    "    hsm.set_segment_actuators(poppy_to_hcipy_index[i], 0, 0, 1.3 * aber_rad_tt / 2)\n",
    "    psm.set_actuator(i, 0, 0, 1.3 * aber_rad_tt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display both segmented mirrors in OPD\n",
    "\n",
    "# HCIPy\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.title('OPD for HCIPy SM')\n",
    "hcipy.imshow_field(hsm.surface * 2, mask=aper, cmap='RdBu_r', vmin=-5e-7, vmax=5e-7)\n",
    "plt.colorbar()\n",
    "\n",
    "# Poppy\n",
    "plt.figure(figsize=(8,8))\n",
    "psm.display(what='opd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Propagate to image plane\n",
    "## HCIPy\n",
    "# Propagate from pupil plane through SM to image plane\n",
    "im_pistoned_hc = prop(hsm(wf))\n",
    "\n",
    "## Poppy\n",
    "# Make an optical system with the Poppy SM and a detector\n",
    "osys = poppy.OpticalSystem()\n",
    "osys.add_pupil(psm)\n",
    "pxscle = 0.0031 * fac  # I'm tweaking pixelscale and fov_arcsec to match the HCIPy image\n",
    "fovarc = 0.05 * fac\n",
    "osys.add_detector(pixelscale=pxscle, fov_arcsec=fovarc, oversample=10)\n",
    "\n",
    "# Calculate the PSF\n",
    "psf = osys.calc_psf(wavelength)\n",
    "\n",
    "# Get the PSF as an array\n",
    "im_pistoned_pop = psf[0].data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Display intensity of both cases image plane\n",
    "plt.figure(figsize=(18, 9))\n",
    "plt.suptitle('Image plane after SM forrandom arangement')\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "hcipy.imshow_field(np.log10(im_pistoned_hc.intensity/norm_hc))\n",
    "plt.title('HCIPy random arangement')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(np.log10(im_pistoned_pop/norm_pop), origin='lower')\n",
    "plt.title('Poppy tipped arangement')\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "level": "intermediate",
  "thumbnail_figure_index": -3
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
